import base64
import random
from io import BytesIO

import numpy as np
from PIL import Image
from flask import Flask, render_template

import data
from train import load_model_weight
import time

app = Flask(__name__)

model = load_model_weight('./model/weights_with_data_enhancement.h5')

test_data, test_label = data.test_data()

categories = [i.decode() for i in data.categories()]
categories_len = len(categories)


@app.route('/', methods=['GET'])
def index():
    """
    创建 "/" 路由，并定义请求方法为GET
    :return: none
    """
    _index = random.randint(0, 10000)
    _img_arr = test_data[_index].copy()

    t1 = int(round(time.time() * 1000))
    _res = model.predict(_img_arr.reshape(-1, 32, 32, 3) / 255.0)
    t2 = int(round(time.time() * 1000))
    use_time = t2 - t1
    _res_category_score = np.max(_res)
    _res_category_index = np.argmax(_res)
    _category = categories[int(_res_category_index)]
    
    _img_arr = test_data[_index].copy().reshape(32, 32, 3)
    _img_arr = np.transpose(_img_arr, (2, 0, 1))
    _image_rgb = (Image.fromarray(_img_arr[0]), Image.fromarray(_img_arr[1]), Image.fromarray(_img_arr[2]))
    _img = Image.merge('RGB', _image_rgb)
    _img = _img.resize((320, 320), Image.ANTIALIAS)
    _buffered = BytesIO()
    _img.save(_buffered, format="PNG")
    _img_base64 = base64.b64encode(_buffered.getvalue()).decode()
    return render_template('index.html', img=_img_base64, category=_category, use_time=use_time,
                           category_real=categories[test_label[_index]], category_score=_res_category_score,
                           all_predict_score=[round(i, 5) for i in _res.tolist()[0]], all_categories=categories)


if __name__ == '__main__':
    app.run(host="127.0.0.1", port=5000)
